#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import sys
sys.path.append("/home/ma-user/work/")
import argparse
import os
# import onnx
import torch

from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import build_detection_test_loader, MetadataCatalog, DatasetCatalog
from detectron2.evaluation import COCOEvaluator, inference_on_dataset, print_csv_format
from detectron2.export import Caffe2Tracer, add_export_config
from detectron2.modeling import build_model
from detectron2.utils.logger import setup_logger

from detectron2.data.datasets import register_coco_instances
import random
from detectron2.utils.visualizer import Visualizer
from detectron2.engine import DefaultPredictor
from fsdet.engine import DefaultTrainer, default_argument_parser, default_setup

import cv2
import torch
from detectron2.data import MetadataCatalog
from detectron2.utils.video_visualizer import VideoVisualizer
from detectron2.utils.visualizer import ColorMode, Visualizer

from fsdet.engine import DefaultTrainer, default_argument_parser, default_setup
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.export import Caffe2Tracer, add_export_config
from detectron2.modeling import build_model
from detectron2.utils.logger import setup_logger

def setup(args):
    cfg = get_cfg()
    cfg.DATALOADER.NUM_WORKERS = 0
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    return cfg

def get_caffe(cfg):
    npu_device = torch.device('npu:0')
    model = build_model(cfg)
    DetectionCheckpointer(model).resume_or_load(cfg.MODEL.WEIGHTS)
    data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0])
    first_batch = next(iter(data_loader))

    # convert and save caffe2 model
    tracer = Caffe2Tracer(cfg, model, first_batch)
    c2_model=tracer.export_caffe2()
    c2_model.save_protobuf(".")
    
if __name__ == "__main__":
    args=default_argument_parser().parse_args()
    cfg=setup(args)
    get_caffe(cfg)
    